#include "stdafx.h"
#include "RegistryCommands.h"
#include "AutoHandle.h"


NtApi ntapi;
extern AutoHandle _autoHandle;


KEY_INFORMATION_CLASS GetKeyClassParameter( const std::wstring& param )
{
	if ( !_wcsicmp( param.c_str(), L"Basic" ) )
	{
		return KeyBasicInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Node" ) )
	{
		return KeyNodeInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Full" ) )
	{
		return KeyFullInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Name" ) )
	{
		return KeyNameInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Cached" ) )
	{
		return KeyCachedInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Flags" ) )
	{
		return KeyFlagsInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Virtualization" ) )
	{
		return KeyVirtualizationInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"HandleTags" ) )
	{
		return KeyHandleTagsInformation;
	}
	return (KEY_INFORMATION_CLASS)-1;
}

KEY_VALUE_INFORMATION_CLASS GetKeyValueClassParameter( const std::wstring& param )
{
	if ( !_wcsicmp( param.c_str(), L"Basic" ) )
	{
		return KeyValueBasicInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Full" ) )
	{
		return KeyValueFullInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"Partial" ) )
	{
		return KeyValuePartialInformation;
	}
	else if ( !_wcsicmp( param.c_str(), L"FullAlign64" ) )
	{
		return KeyValueFullInformationAlign64;
	}
	else if ( !_wcsicmp( param.c_str(), L"PartialAlign64" ) )
	{
		return KeyValuePartialInformationAlign64;
	}
	return (KEY_VALUE_INFORMATION_CLASS)-1;
}

ULONG GetRegDataType( const std::wstring& param )
{
	if ( !_wcsicmp( param.c_str(), L"DWORD" ) )
	{
		return REG_DWORD;
	}
	else if ( !_wcsicmp( param.c_str(), L"SZ" ) )
	{
		return REG_SZ;
	}
	else if ( !_wcsicmp( param.c_str(), L"EXPAND_SZ" ) )
	{
		return REG_EXPAND_SZ;
	}
	else if ( !_wcsicmp( param.c_str(), L"BINARY" ) )
	{
		return REG_BINARY;
	}
	else if ( !_wcsicmp( param.c_str(), L"LINK" ) )
	{
		return REG_LINK;
	}
	return REG_DWORD;
}

CloseCommand::CloseCommand(void)
	: BaseCommand( L"Close" )
{
}

CloseCommand::~CloseCommand(void)
{
}

void CloseCommand::DisplayHelp()
{
	wprintf( L"Help: close -handle [handle]\n" );
}

bool CloseCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 1 ||
		params.find( L"handle" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	HANDLE handle = GetHandleParameter(params, L"handle");
	_autoHandle.EraseHandle( handle );
	NTSTATUS status = ntapi.NtClose( handle );
	wprintf( L"Status = 0x%08x\n", status );
	return false;
}

CreateKeyCommand::CreateKeyCommand(void)
	: BaseCommand( L"CreateKey" )
{
}

CreateKeyCommand::~CreateKeyCommand(void)
{
}

void CreateKeyCommand::DisplayHelp()
{
	wprintf( L"Help: createkey [-handle [root_handle]] -name [key_name] [-access [desired_access]] [-options [key_options]]\n" );
}

bool CreateKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() < 1 ||
		params.find( L"name" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	std::wstring keyName = GetNamedParameter(params, L"name");
	UNICODE_STRING KeyName = RTL_CONSTANT_STRING( const_cast<wchar_t*>(keyName.c_str()) );
	OBJECT_ATTRIBUTES Oa = RTL_CONSTANT_OBJECT_ATTRIBUTES( GetHandleParameter(params, L"handle", 0), &KeyName, OBJ_CASE_INSENSITIVE );
	ULONG Disposition;
	HANDLE NewHandle;
	NTSTATUS status = ntapi.NtCreateKey( &NewHandle, GetAccessParameter(params, L"access", L"KR"), &Oa, 0, NULL, GetNamedParameterInt(params, L"options", 0), &Disposition );
	wprintf( L"Status = 0x%08x\n", status );
	if ( NT_SUCCESS( status ) )
	{
		wprintf( L"Disposition = %s\n", Disposition == REG_CREATED_NEW_KEY ? L"Created" : L"Opened" );
		long autoIndex = _autoHandle.PushHandle( NewHandle );
		wprintf( L"Handle = %d (AUTO-%d)\n", NewHandle, autoIndex );
	}
	return false;
}

DeleteKeyCommand::DeleteKeyCommand(void)
	: BaseCommand( L"DeleteKey" )
{
}

DeleteKeyCommand::~DeleteKeyCommand(void)
{
}

void DeleteKeyCommand::DisplayHelp()
{
	wprintf( L"Help: deletekey -handle [handle]\n" );
}

bool DeleteKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 1 ||
		params.find( L"handle" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	NTSTATUS status = ntapi.NtDeleteKey( GetHandleParameter(params, L"handle") );
	wprintf( L"Status = 0x%08x\n", status );
	return false;
}

DeleteValueKeyCommand::DeleteValueKeyCommand(void)
	: BaseCommand( L"DeleteValueKey" )
{
}

DeleteValueKeyCommand::~DeleteValueKeyCommand(void)
{
}

void DeleteValueKeyCommand::DisplayHelp()
{
	wprintf( L"Help: deletevaluekey -handle [handle] -name [value_name]\n" );
}

bool DeleteValueKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 2 ||
		params.find( L"handle" ) == params.end() ||
		params.find( L"name" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	std::wstring valueName = GetNamedParameter(params, L"name");
	UNICODE_STRING ValueName = RTL_CONSTANT_STRING( const_cast<wchar_t*>(valueName.c_str()) );
	NTSTATUS status = ntapi.NtDeleteValueKey( GetHandleParameter(params, L"handle"), &ValueName );
	wprintf( L"Status = 0x%08x\n", status );
	return false;
}

EnumKeyCommand::EnumKeyCommand(void)
	: BaseCommand( L"EnumKey" )
{
}

EnumKeyCommand::~EnumKeyCommand(void)
{
}

void EnumKeyCommand::DisplayHelp()
{
	wprintf( L"Help: enumkey -handle [handle] -index [index] -class [information_class] -bufferlen [buffer_size]\n" );
}

bool EnumKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 4 ||
		params.find( L"handle" ) == params.end() ||
		params.find( L"index" ) == params.end() ||
		params.find( L"class" ) == params.end() ||
		params.find( L"bufferlen" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	unsigned long Index = GetNamedParameterInt(params, L"index");
	unsigned long cbBuffer = GetNamedParameterInt(params, L"bufferlen");
	BYTE* pBuffer = (BYTE*) malloc( cbBuffer );
	unsigned long cbResult;
	NTSTATUS status = ntapi.NtEnumerateKey( GetHandleParameter(params, L"handle"), Index, GetKeyClassParameter(GetNamedParameter(params, L"class")), pBuffer, cbBuffer, &cbResult );
	wprintf( L"Status = 0x%08x\n", status );
	wprintf( L"ResultLength = %d\n", cbResult );
	if ( NT_SUCCESS( status ) )
	{
		dumpBuffer( pBuffer, cbResult );
	}
	free( pBuffer );
	return false;
}

EnumValueKeyCommand::EnumValueKeyCommand(void)
	: BaseCommand( L"EnumValueKey" )
{
}

EnumValueKeyCommand::~EnumValueKeyCommand(void)
{
}

void EnumValueKeyCommand::DisplayHelp()
{
	wprintf( L"Help: enumvaluekey -handle [handle] -index [index] -class [information_class] -bufferlen [buffer_size]\n" );
}

bool EnumValueKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 4 ||
		params.find( L"handle" ) == params.end() ||
		params.find( L"index" ) == params.end() ||
		params.find( L"class" ) == params.end() ||
		params.find( L"bufferlen" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	unsigned long Index = GetNamedParameterInt(params, L"index");
	unsigned long cbBuffer = GetNamedParameterInt(params, L"bufferlen");
	BYTE* pBuffer = (BYTE*) malloc( cbBuffer );
	unsigned long cbResult;
	NTSTATUS status = ntapi.NtEnumerateValueKey( GetHandleParameter(params, L"handle"), Index, GetKeyValueClassParameter(GetNamedParameter(params, L"class")), pBuffer, cbBuffer, &cbResult );
	wprintf( L"Status = 0x%08x\n", status );
	wprintf( L"ResultLength = %d\n", cbResult );
	if ( NT_SUCCESS( status ) )
	{
		dumpBuffer( pBuffer, cbResult );
	}
	free( pBuffer );
	return false;
}

OpenKeyCommand::OpenKeyCommand(void)
	: BaseCommand( L"OpenKey" )
{
}

OpenKeyCommand::~OpenKeyCommand(void)
{
}

void OpenKeyCommand::DisplayHelp()
{
	wprintf( L"Help: openkey [-handle [root_handle]] -name [key_name] [-access [desired_access]]\n" );
}

bool OpenKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() < 1 ||
		params.find( L"name" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	std::wstring keyName = GetNamedParameter(params, L"name");
	UNICODE_STRING KeyName = RTL_CONSTANT_STRING( const_cast<wchar_t*>(keyName.c_str()) );
	OBJECT_ATTRIBUTES Oa = RTL_CONSTANT_OBJECT_ATTRIBUTES( GetHandleParameter(params, L"handle", 0), &KeyName, OBJ_CASE_INSENSITIVE );
	HANDLE NewHandle;
	NTSTATUS status = ntapi.NtOpenKey( &NewHandle, GetAccessParameter(params, L"access", L"KR"), &Oa );
	wprintf( L"Status = 0x%08x\n", status );
	if ( NT_SUCCESS( status ) )
	{
		long autoIndex = _autoHandle.PushHandle( NewHandle );
		wprintf( L"Handle = %d (AUTO-%d)\n", NewHandle, autoIndex );
	}
	return false;
}

OpenKeyExCommand::OpenKeyExCommand(void)
	: BaseCommand( L"OpenKeyEx" )
{
}

OpenKeyExCommand::~OpenKeyExCommand(void)
{
}

void OpenKeyExCommand::DisplayHelp()
{
	wprintf( L"Help: openkeyex [-handle [root_handle]] -name [key_name] [-access [desired_access]] [-options [key_options]]\n" );
}

bool OpenKeyExCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() < 1 ||
		params.find( L"name" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	std::wstring keyName = GetNamedParameter(params, L"name");
	UNICODE_STRING KeyName = RTL_CONSTANT_STRING( const_cast<wchar_t*>(keyName.c_str()) );
	ULONG Options = GetNamedParameterInt(params, L"options", 0);
	ULONG OaFlags = OBJ_CASE_INSENSITIVE;
	if ( Options & REG_OPTION_OPEN_LINK ) OaFlags |= OBJ_OPENLINK;
	OBJECT_ATTRIBUTES Oa = RTL_CONSTANT_OBJECT_ATTRIBUTES( GetHandleParameter(params, L"handle", 0), &KeyName, OaFlags );
	HANDLE NewHandle;
	NTSTATUS status = ntapi.NtOpenKeyEx( &NewHandle, GetAccessParameter(params, L"access", L"KR"), &Oa, Options );
	wprintf( L"Status = 0x%08x\n", status );
	if ( NT_SUCCESS( status ) )
	{
		long autoIndex = _autoHandle.PushHandle( NewHandle );
		wprintf( L"Handle = %d (AUTO-%d)\n", NewHandle, autoIndex );
	}
	return false;
}

QueryKeyCommand::QueryKeyCommand(void)
	: BaseCommand( L"QueryKey" )
{
}

QueryKeyCommand::~QueryKeyCommand(void)
{
}

void QueryKeyCommand::DisplayHelp()
{
	wprintf( L"Help: querykey -handle [handle] -class [information_class] -bufferlen [buffer_size]\n" );
}

bool QueryKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 3 ||
		params.find( L"handle" ) == params.end() ||
		params.find( L"class" ) == params.end() ||
		params.find( L"bufferlen" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	unsigned long cbBuffer = GetNamedParameterInt(params, L"bufferlen");
	BYTE* pBuffer = (BYTE*) malloc( cbBuffer );
	unsigned long cbResult;
	NTSTATUS status = ntapi.NtQueryKey( GetHandleParameter(params, L"handle"), GetKeyClassParameter(GetNamedParameter(params, L"class")), pBuffer, cbBuffer, &cbResult );
	wprintf( L"Status = 0x%08x\n", status );
	wprintf( L"ResultLength = %d\n", cbResult );
	if ( NT_SUCCESS( status ) )
	{
		dumpBuffer( pBuffer, cbResult );
	}
	free( pBuffer );
	return false;
}

QueryValueKeyCommand::QueryValueKeyCommand(void)
	: BaseCommand( L"QueryValueKey" )
{
}

QueryValueKeyCommand::~QueryValueKeyCommand(void)
{
}

void QueryValueKeyCommand::DisplayHelp()
{
	wprintf( L"Help: queryvaluekey -handle [handle] -name [value_name] -class [information_class] -bufferlen [buffer_size]\n" );
}

bool QueryValueKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 4 ||
		params.find( L"handle" ) == params.end() ||
		params.find( L"name" ) == params.end() ||
		params.find( L"class" ) == params.end() ||
		params.find( L"bufferlen" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	unsigned long cbBuffer = GetNamedParameterInt(params, L"bufferlen");
	BYTE* pBuffer = (BYTE*) malloc( cbBuffer );
	unsigned long cbResult;
	std::wstring valueName = GetNamedParameter(params, L"name");
	UNICODE_STRING ValueName = RTL_CONSTANT_STRING( const_cast<wchar_t*>(valueName.c_str()) );
	NTSTATUS status = ntapi.NtQueryValueKey( GetHandleParameter(params, L"handle"), &ValueName, GetKeyValueClassParameter(GetNamedParameter(params, L"class")), pBuffer, cbBuffer, &cbResult );
	wprintf( L"Status = 0x%08x\n", status );
	wprintf( L"ResultLength = %d\n", cbResult );
	if ( NT_SUCCESS( status ) )
	{
		dumpBuffer( pBuffer, cbResult );
	}
	free( pBuffer );
	return false;
}

SetValueKeyCommand::SetValueKeyCommand(void)
	: BaseCommand( L"SetValueKey" )
{
}

SetValueKeyCommand::~SetValueKeyCommand(void)
{
}

void SetValueKeyCommand::DisplayHelp()
{
	wprintf( L"Help: SetValueKey -handle [handle] -name [value_name] -type [reg_type] -data [data]\n" );
}

bool SetValueKeyCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 4 ||
		params.find( L"handle" ) == params.end() ||
		params.find( L"name" ) == params.end() ||
		params.find( L"type" ) == params.end() ||
		params.find( L"data" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	BYTE* pBuffer = NULL;
	unsigned long cbBuffer;

	//	Convert string from escaped data to registry buffer
	switch ( GetRegDataType(GetNamedParameter(params, L"type")) )
	{
	case REG_DWORD:
		{
			unsigned long dataNumber = GetNamedParameterInt(params, L"data");
			cbBuffer = sizeof(unsigned long);
			pBuffer = (BYTE*) malloc( cbBuffer );
			*(unsigned long*)pBuffer = dataNumber;
		}
		break;

	case REG_SZ:
	case REG_EXPAND_SZ:
	case REG_LINK:
		{
			std::wstring dataString = GetNamedParameter(params, L"data");
			cbBuffer = (unsigned long)dataString.size() * sizeof(wchar_t);
			pBuffer = (BYTE*) malloc( cbBuffer );
			wchar_t* pStringBuffer = (wchar_t*)pBuffer;
			memset( pBuffer, 0, cbBuffer );
			size_t idx = 0;
			for ( size_t i = 0; i < dataString.size(); i++ )
			{
				if ( dataString[i] == L'\\' )
				{
					switch (dataString[i+1])
					{
					case L'\\':
						pStringBuffer[idx] = L'\\';
						i++;
						cbBuffer -= sizeof(wchar_t);
						break;
					case L'0':
						pStringBuffer[idx] = L'\0';
						i++;
						cbBuffer -= sizeof(wchar_t);
						break;
					case L'x':
						{
							wchar_t byteString[3] = { dataString[i+2], dataString[i+3], 0 };
							wchar_t* pEndScan;
							pStringBuffer[idx] = (BYTE)wcstoul( byteString, &pEndScan, 16 );
							i += 2;
							cbBuffer -= (3 * sizeof(wchar_t));
						}
						break;
					default:
						//	Unrecognized escape character - ignore the slash
						i++;
						cbBuffer -= sizeof(wchar_t);
						break;
					}
				}
				else
				{
					pStringBuffer[idx] = dataString[i];
				}
				idx++;
			}
		}
		break;

	case REG_BINARY:
		{
			std::wstring dataString = GetNamedParameter(params, L"data");
			cbBuffer = (unsigned long)dataString.size();
			pBuffer = (BYTE*) malloc( cbBuffer );
			memset( pBuffer, 0, cbBuffer );
			cbBuffer = 0;
			size_t idx = 0;
			for ( size_t i = 0; i < dataString.size(); i++ )
			{
				if ( dataString[i] == L'\\' )
				{
					switch (dataString[i+1])
					{
					case L'\\':
						pBuffer[idx] = '\\';
						i++;
						cbBuffer++;
						break;
					case L'0':
						pBuffer[idx] = '\0';
						i++;
						cbBuffer++;
						break;
					case L'x':
						{
							wchar_t byteString[3] = { dataString[i+2], dataString[i+3], 0 };
							wchar_t* pEndScan;
							pBuffer[idx] = (BYTE)wcstoul( byteString, &pEndScan, 16 );
							i += 3;
							cbBuffer++;
						}
						break;
					default:
						//	Unrecognized escape character - ignore the slash
						break;
					}
				}
				else
				{
					pBuffer[idx] = (BYTE)dataString[i];
					cbBuffer++;
				}
				idx++;
			}
		}
		break;
	default:
		wprintf( L"Unknown registry data type: try DWORD, SZ, EXPAND_SZ, LINK, or BINARY\n" );
		break;
	}

	if ( pBuffer )
	{
		std::wstring valueName = GetNamedParameter(params, L"name");
		UNICODE_STRING ValueName = RTL_CONSTANT_STRING( const_cast<wchar_t*>(valueName.c_str()) );
		NTSTATUS status = ntapi.NtSetValueKey( GetHandleParameter(params, L"handle"), &ValueName, 0, GetRegDataType(GetNamedParameter(params, L"type")), pBuffer, cbBuffer );
		wprintf( L"Status = 0x%08x\n", status );
		free( pBuffer );
	}
	return false;
}
